#include "vtkMeshExport.hpp"
#include "vtkMeshImport.hpp"
#include "vtkFieldExport_boundary.hpp"
#include "vtkFieldImport_boundary.hpp"
#include "sweSolver.hpp"

#include "bc_swe_all.hpp"

#include "meshTypeFunctions.hpp"
#include "bc_weir_funcs.hpp"

#include <catch2/catch_test_macros.hpp>
#include <catch2/catch_approx.hpp>

#include <cmath>

using namespace hms;
using namespace hms::vtk;
namespace xml = tinyxml2;
using Catch::Approx;

TEST_CASE("Read boundary conditions"){


	auto test = []( auto&& meshDummy ){
		using MeshType = remove_qualifiers<decltype(meshDummy)>;

		MeshVariant meshVar { uniMeshAs<MeshType>(3,4,1,1,true) };

		const MeshType& mesh { std::get<MeshType>(meshVar) };

		Time time {10};
		time.writeTimeInterval(1);

		SWESolver<MeshType> solver {mesh, time};

		solver.v().values().colwise() = Array2s{1, 0};
		solver.d().values() = 1;
		solver.z().values() = 0.0;

		/* set all BCs to ZeroGradient, just so that something is being written
		 * on export */
		auto allBCsToZeroGradient = [&](){
			for ( const auto& patch : mesh.boundaryPatches() ){
				for ( Field<MeshType>& field : solver.fields() ){
					using ZG = ZeroGradient<MeshType>;
					field.template setBoundaryCondition<ZG>(patch);
				}
			}
		};
		allBCsToZeroGradient();

		std::unique_ptr<xml::XMLDocument> boundaryDoc { boundaryToXML(meshVar) };
		writeBoundaryConditions( boundaryDoc.get(), solver );

		const BoundaryPatch
			& top { mesh.boundaryPatch("top") },
			& left { mesh.boundaryPatch("left") },
			& right { mesh.boundaryPatch("right") },
			& bottom { mesh.boundaryPatch("bottom") };

		xml::XMLElement* bElem { getBoundaryNode( boundaryDoc.get() ) };

		auto getBC = [&]( const BoundaryPatch& patch, const char* field )
			-> xml::XMLElement*
		{
			return getChildNodeByAttributeValue(
				getBoundaryPatch(bElem, patch), "BoundaryCondition", "Field",
				field
			);
		};
		
		auto reimportCheck = [&](){
			/* boundary patches, without boundary conditions */
			std::unique_ptr<xml::XMLDocument>
				boundaryDoc1 { boundaryToXML(meshVar) },
				boundaryDoc2 { boundaryToXML(meshVar) };
			std::stringstream s1, s2;

			/* write conditions to first doc, but not to second,
			 * then store the doc's state as a string stream */
			writeBoundaryConditions( boundaryDoc1.get(), solver );
			toStream( boundaryDoc1.get(), s1 );

			/* reset BCs of all fields, so that the info is *only*
			 * in doc1 and the string stream */
			allBCsToZeroGradient();

			/* read the doc in, to restore the boundary conditions for all fields,
			 * then write them to the second string stream */
			readBoundaryConditions( boundaryDoc1.get(), solver, false );
			writeBoundaryConditions( boundaryDoc2.get(), solver );
			toStream( boundaryDoc2.get(), s2 );

			/* now the content of both strings are the same,
			 * iff export and import are congruent */
			REQUIRE( s1.str() == s2.str() );
		};

		SECTION("StaticValue"){
			auto bcVelocity = getBC(top, "Velocity");
			auto bcElev = getBC(top, "BottomElevation");

			bcVelocity->SetAttribute("Type", "staticValue");
			bcVelocity->SetAttribute("isUniform", true);
			bcVelocity->SetAttribute("isFunction", false);
			bcElev->SetAttribute("Type", "staticValue");
			bcElev->SetAttribute("isUniform", true);
			bcElev->SetAttribute("isFunction", false);

			bcVelocity->SetText("2 3");
			bcElev->SetText("-0.3");

			readBoundaryConditions( boundaryDoc.get(), solver, false );

			REQUIRE( solver.v().ghostValues(top).isApprox(
				Array2s{2, 3}.replicate(1,3)
			) );
			REQUIRE( solver.z().ghostValues(top).isApproxToConstant(
				-0.3
			));

			bcVelocity = getBC( bottom, "Velocity" );
			bcElev = getBC( bottom, "BottomElevation" );

			bcVelocity->SetAttribute("Type", "staticValue");
			bcVelocity->SetAttribute("isUniform", false);
			bcVelocity->SetAttribute("isFunction", false);
			bcElev->SetAttribute("Type", "staticValue");
			bcElev->SetAttribute("isUniform", false);
			bcElev->SetAttribute("isFunction", false);

			bcVelocity->SetText("1 2 3 4 5 6");
			bcElev->SetText("9 8 7");

			readBoundaryConditions( boundaryDoc.get(), solver, false );

			REQUIRE( solver.v().ghostValues(bottom).isApprox(
				ArrayNs<6>{1,2,3,4,5,6}.reshaped(2,3)
			) );
			REQUIRE( solver.z().ghostValues(bottom).isApprox(
				RowArray3s{9, 8, 7}
			));

			bcVelocity = getBC( left, "Velocity" );
			bcElev = getBC( left, "BottomElevation" );

			bcVelocity->SetAttribute("Type", "staticValue");
			bcVelocity->SetAttribute("isUniform", false);
			bcVelocity->SetAttribute("isFunction", true);
			bcElev->SetAttribute("Type", "staticValue");
			bcElev->SetAttribute("isUniform", false);
			bcElev->SetAttribute("isFunction", true);

			writeFunction(bcVelocity, "x+y", 0);
			writeFunction(bcVelocity, "x*y", 1);
			writeFunction(bcElev, "x-y", 0);

			readBoundaryConditions( boundaryDoc.get(), solver, false );

			REQUIRE( solver.v().ghostValues(left).isApprox(
				ArrayNs<8>{ 0, -0.25, 1, -0.75, 2, -1.25, 3, -1.75 }.reshaped(2,4)
			) );
			REQUIRE( solver.z().ghostValues(left).isApprox(
				RowArray4s{ -1, -2, -3, -4 }
			) );

			bcVelocity = getBC( right, "Velocity" );
			bcElev = getBC( right, "BottomElevation" );

			bcVelocity->SetAttribute("Type", "staticValue");
			bcVelocity->SetAttribute("isUniform", true);
			bcVelocity->SetAttribute("isFunction", true);
			bcElev->SetAttribute("Type", "staticValue");
			bcElev->SetAttribute("isUniform", true);
			bcElev->SetAttribute("isFunction", true);

			writeFunction(bcVelocity, "pi", 0);
			writeFunction(bcVelocity, "g", 1);
			writeFunction(bcElev, "e", 0);

			std::cerr.setstate(std::ios_base::failbit);
			readBoundaryConditions( boundaryDoc.get(), solver, false );
			std::cerr.clear();

			REQUIRE( solver.v().ghostValues(right).isApprox(
				Array2s{ math::pi, phys::g }.replicate(1,4)
			) );
			REQUIRE( solver.z().ghostValues(right).isApproxToConstant(
				math::e
			) );

			SECTION( "Skipping z field"){
				solver.z().ghostValues(top) = 16;
				std::cerr.setstate(std::ios_base::failbit);
				readBoundaryConditions( boundaryDoc.get(), solver, true );
				std::cerr.clear();
				
				REQUIRE( solver.z().ghostValues(top).isApproxToConstant(16) );
			}
		}

		SECTION("DynamicValue"){
			using DV = DynamicValue<MeshType>;
			auto getDynVal = [](
				const Field<MeshType>& field,
				const BoundaryPatch& patch
			) -> const DV* {
				auto dv { dynamic_cast<const DV*>(
					&(field.boundaryCondition(patch))
				)};
				if (!dv)
					assert(false);
				return dv;
			};

			auto compareTimeSeries = [&getDynVal](
				const Field<MeshType>& field,
				const BoundaryPatch& patch,
				const TimeSeries& target
			){
				auto dv = getDynVal(field, patch);

				const TimeSeries& actual {
					dv->values()
				};
				REQUIRE(target.size() == actual.size());
				for ( size_t i{0}; i < target.size(); ++i){
					REQUIRE( approxEqual(target[i].first, actual[i].first) );
					// std::cout
						// <<   "target:\n" << target[i].second
						// << "\nactual:\n" << actual[i].second << "\n";
					REQUIRE( target[i].second.isApprox(actual[i].second) );
				}
			};

			auto compareFunctionStrings = [&getDynVal](
				const Field<MeshType>& field,
				const BoundaryPatch& patch,
				const std::vector<std::string>& target
			){
				auto dv = getDynVal(field, patch);
				const std::vector<FunctionParserWrapper>& actual {
					dv->functions()
				};
				REQUIRE(target.size() == actual.size());
				for ( size_t i{0}; i < target.size(); ++i){
					REQUIRE( actual[i].functionString() == target[i] );
				}
			};

			SECTION("Uniform time series"){
				auto bcVelocity = getBC(top, "Velocity");
				auto bcElev = getBC(top, "BottomElevation");

				bcVelocity->SetAttribute("Type", "dynamicValue");
				bcVelocity->SetAttribute("isUniform", true);
				bcVelocity->SetAttribute("isFunction", false);
				bcVelocity->SetAttribute("NumberOfTimeValues", 2);
				bcElev->SetAttribute("Type", "dynamicValue");
				bcElev->SetAttribute("isUniform", true);
				bcElev->SetAttribute("isFunction", false);
				bcElev->SetAttribute("NumberOfTimeValues", 2);

				TimeSeries
					vSeries {
						{0, Array2s{-1, 3}},
						{5, Array2s{-2, 6}}
					},
					zSeries {
						{1, ArrayNs<1>{-2}},
						{4, ArrayNs<1>{-3}}
					};

				writeTimeSeries(bcVelocity, vSeries);
				writeTimeSeries(bcElev, zSeries);

				std::cerr.setstate(std::ios_base::failbit);
				readBoundaryConditions( boundaryDoc.get(), solver, false );
				std::cerr.clear();

				compareTimeSeries( solver.v(), top, vSeries );
				compareTimeSeries( solver.z(), top, zSeries );
			}
			SECTION("Non-uniform time series"){

				auto bcVelocity = getBC(bottom, "Velocity");
				auto bcElev = getBC(bottom, "BottomElevation");

				bcVelocity->SetAttribute("Type", "dynamicValue");
				bcVelocity->SetAttribute("isUniform", false);
				bcVelocity->SetAttribute("isFunction", false);
				bcVelocity->SetAttribute("NumberOfTimeValues", 2);
				bcElev->SetAttribute("Type", "dynamicValue");
				bcElev->SetAttribute("isUniform", false);
				bcElev->SetAttribute("isFunction", false);
				bcElev->SetAttribute("NumberOfTimeValues", 2);
			
				ArrayXXs v1 (2,3), v2;
				v2.resizeLike(v1);
				v1.setRandom();
				v2.setRandom();

				TimeSeries
					vSeries {
						{0, std::move(v1)},
						{5, std::move(v2)}
					},
					zSeries {
						{0, RowArray3s{3,4,5}},
						{4, RowArray3s{8,7,6}}
					};

				writeTimeSeries(bcVelocity, vSeries);
				writeTimeSeries(bcElev, zSeries);

				readBoundaryConditions( boundaryDoc.get(), solver, false );

				compareTimeSeries( solver.v(), bottom, vSeries );
				compareTimeSeries( solver.z(), bottom, zSeries );
			}
			SECTION("Function"){

				auto bcVelocity = getBC(bottom, "Velocity");
				auto bcElev = getBC(bottom, "BottomElevation");

				bcVelocity->SetAttribute("Type", "dynamicValue");
				bcVelocity->SetAttribute("isUniform", true);
				bcVelocity->SetAttribute("isFunction", true);
				bcElev->SetAttribute("Type", "dynamicValue");
				bcElev->SetAttribute("isUniform", true);
				bcElev->SetAttribute("isFunction", true);

				std::string
					v0 {"1/2 * g*t"},
					v1 {"3*pi*t"},
					z0 {"5*e"};
				writeFunction(bcVelocity, v0, 0);
				writeFunction(bcVelocity, v1, 1);
				writeFunction(bcElev, z0, 0);

				readBoundaryConditions( boundaryDoc.get(), solver, false );

				compareFunctionStrings( solver.v(), bottom, {v0, v1} );
				compareFunctionStrings( solver.z(), bottom, {z0} );
			}
		}
		SECTION( "SWE boundaries" ){
			SECTION( "Weir" ){
				xml::XMLElement
					* bcWaterDepth = getBC(top, "WaterDepth"),
					* bcVelocity   = getBC(top, "Velocity"),
					* bcBottomElev = getBC(top, "BottomElevation");

				bcWaterDepth->SetAttribute("Type", "weir");
				bcVelocity  ->SetAttribute("Type", "weir");
				bcBottomElev->SetAttribute("Type", "weir");

				scalar
					alpha {30},
					alphaRad { degToRad(alpha) },
					cp { bc_weir::pressureFactor(alphaRad) };

				bcVelocity  ->SetAttribute("angle", alpha);
				bcWaterDepth->SetAttribute("angle", alpha);
				/* to check conversion, use radians of same angle here */
				bcBottomElev->SetAttribute( "angle", alphaRad );
				bcBottomElev->SetAttribute("isRadians", true);

				readBoundaryConditions( boundaryDoc.get(), solver, false );

				using BCType = boundaryCondition::Type;

				auto checkType = [&](
					const Field<MeshType>& field, BCType bcType
				){
					REQUIRE( field.boundaryCondition(top).type() == bcType );
				};
				checkType( solver.v(), BCType::weir_velocity );
				checkType( solver.d(), BCType::weir_waterdepth );
				checkType( solver.z(), BCType::weir_bottomElevation );

				auto checkAngle = [&](
					const Field<MeshType>& field, auto&& dummyType
				){
					using BC = remove_qualifiers<decltype(dummyType)>;
					const BC* bc { dynamic_cast<const BC*>(
						&( field.boundaryCondition(top) )
					) };
					REQUIRE(bc);
					REQUIRE( bc->alpha() == Approx( alphaRad ) );
					REQUIRE( bc->cp()    == Approx( cp ) );
				};

				checkAngle( solver.v(), Weir_velocity<MeshType>{} );
				checkAngle( solver.d(), Weir_waterdepth<MeshType>{} );
				checkAngle( solver.z(), Weir_bottomElevation<MeshType>{} );

				reimportCheck();
			}
			SECTION( "Orthogonal discharge" ){
				xml::XMLElement
					* bcWaterDepth = getBC(top, "WaterDepth"),
					* bcVelocity   = getBC(top, "Velocity");

				bcWaterDepth->SetAttribute("Type", "orthogonalDischarge");
				bcVelocity  ->SetAttribute("Type", "orthogonalDischarge");

				scalar discharge {5};

				bcVelocity->SetAttribute("discharge", discharge);

				readBoundaryConditions( boundaryDoc.get(), solver, false );

				using BCType = boundaryCondition::Type;

				auto checkType = [&](
					const Field<MeshType>& field, BCType bcType
				){
					REQUIRE( field.boundaryCondition(top).type() == bcType );
				};
				checkType( solver.v(), BCType::orthogonalDischarge_velocity );
				checkType( solver.d(), BCType::orthogonalDischarge_waterdepth );

				using BC = OrthogonalDischarge_velocity<MeshType>;

				const BC* bc { dynamic_cast<const BC*>(
					&( solver.v().boundaryCondition(top) )
				) };
				REQUIRE(bc);
				REQUIRE( bc->qBoundary() == Approx(discharge) );

				reimportCheck();
			}
		}
		SECTION( "Exporting and reimporting" ){
			bool uniform{true}, nonUniform{false};
			using DV = DynamicValue<MeshType>;
			using SV = StaticValue<MeshType>;
			using ZG = ZeroGradient<MeshType>;
			using FO_v = FreeOutflow_velocity<MeshType>;
			using FO_w = FreeOutflow_waterdepth<MeshType>;
			using OFV_v = OrthogonalFlowVelocity_velocity<MeshType>;
			using OFV_w = OrthogonalFlowVelocity_waterdepth<MeshType>;

			DV
				vLeft {
					solver.v(), left, solver.time(),
					TimeSeries{
						{0, Array2s{1.5, 0} },
						{1, Array2s{0, -1.5} }
					},
					uniform
				},
				dLeft {
					solver.d(), left, solver.time(),
					TimeSeries{
						{0, Array4s{1,2,3,4}.transpose() },
						{5, Array4s{4,3,2,1}.transpose() }
					},
					nonUniform
				},
				zLeft {
					solver.z(), left, solver.time(),
					std::vector<FunctionParserWrapper>{
						{"sin(t)"}
					},
					uniform,
					solver.z()
				},
				vRight {
					solver.v(), right, solver.time(),
					std::vector<FunctionParserWrapper>{
						{"cos(t)"},
						{"x+2*y"}
					},
					nonUniform,
					solver.z()
				};
				

			solver.v().template setBoundaryCondition<SV>( bottom,
				Array2s{1,2}, uniform );
			solver.d().template setBoundaryCondition<SV>( bottom,
				ArrayNNs<1,3>{0.0, -0.1, 0.1}, nonUniform );
			solver.z().template setBoundaryCondition<ZG>( bottom );

			solver.v().setBoundaryCondition( std::move(vLeft) );
			solver.d().setBoundaryCondition( std::move(dLeft) );
			solver.z().setBoundaryCondition( std::move(zLeft) );

			solver.v().setBoundaryCondition( std::move(vRight) );
			solver.d().template setBoundaryCondition<FO_w>( right, true, 42 );
			solver.z().template setBoundaryCondition<ZG>( right );

			std::cerr.setstate(std::ios_base::failbit);
			solver.v().template setBoundaryCondition<OFV_v>( top, 5 );
			solver.d().template setBoundaryCondition<OFV_w>( top, solver.v(), 5 );
			std::cerr.clear();
			solver.z().template setBoundaryCondition<ZG>( top );

			reimportCheck();
		}
	};

	SECTION("Mesh"){
		test( Mesh{} );
	}
	SECTION("UniMesh"){
		test( UniMesh{} );
	}
	SECTION("RectMesh"){
		test( RectMesh{} );
	}
	SECTION("StructMesh"){
		test( StructMesh{} );
	}
}